#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "mhvi_search.h"

void parse_config(const char *filename, Config *config) {
    FILE *fp = fopen(filename, "r");
    if (!fp) {
        perror("Failed to open file");
        exit(EXIT_FAILURE);
    }

    char line[MAX_LINE_LENGTH];
    while (fgets(line, sizeof(line), fp)) {
        char key[64];
        char value[256];
        if (sscanf(line, "%[^=]=%s", key, value) == 2) {
            if (strcmp(key, "dataset_path") == 0) {
                strncpy(config->dataset_path, value, sizeof(config->dataset_path) - 1);
            } else if (strcmp(key, "query_path") == 0) {
                strncpy(config->query_path, value, sizeof(config->query_path) - 1);
            } else if (strcmp(key, "gt_path") == 0) {
                strncpy(config->gt_path, value, sizeof(config->gt_path) - 1);
            } else if (strcmp(key, "cluster_center_path") == 0) {
                strncpy(config->cluster_center_path, value, sizeof(config->cluster_center_path) - 1);
            } else if (strcmp(key, "offset_list_path") == 0) {
                strncpy(config->offset_list_path, value, sizeof(config->offset_list_path) - 1);
            } else if (strcmp(key, "last_layer_path") == 0) {
                strncpy(config->last_layer_path, value, sizeof(config->last_layer_path) - 1);
            } else if (strcmp(key, "cluster_count") == 0) {
                config->cluster_count = atoi(value);
            } else if (strcmp(key, "train_ratio") == 0) {
                config->train_ratio = atof(value);
            } else if (strcmp(key, "k") == 0) {
                config->k = atoi(value);
            } else if (strcmp(key, "nprobe") == 0) {
                config->nprobe = atoi(value);
            } else if (strcmp(key, "R") == 0) {
                config->R = atoi(value);
            } else if (strcmp(key, "L") == 0) {
                config->L = atoi(value);
            } else if (strcmp(key, "B") == 0) {
                config->B = atoi(value);
            } else if (strcmp(key, "M") == 0) {
                config->M = atoi(value);
            } else if (strcmp(key, "vector_count") == 0) {
                config->n = atoi(value);
            } else if (strcmp(key, "dim") == 0) {
                config->dim = atoi(value);
            } else if (strcmp(key, "max_query") == 0) {
                config->max_query = atoi(value);
            } else if (strcmp(key, "thread_count") == 0) {
                config->thread_count = atoi(value);
            } else if (strcmp(key, "recall_k") == 0) {
                config->recall_k = atoi(value);
            } else if (strcmp(key, "real_k") == 0) {
                config->real_k = atoi(value);
            } else if (strcmp(key, "last_layer_nprobe") == 0) {
                config->last_layer_nprobe = atoi(value);
            }
        }
    }

    fclose(fp);
}

void print_config(const Config *config) {
    printf("==================Config==================\n");
    printf("dataset_path: %s\n", config->dataset_path);
    printf("query_path: %s\n", config->query_path);
    printf("gt_path: %s\n", config->gt_path);
    printf("cluster_center_path: %s\n", config->cluster_center_path);
    printf("offset_list_path: %s\n", config->offset_list_path);
    printf("last_layer_path: %s\n", config->last_layer_path);
    printf("cluster_count: %d\n", config->cluster_count);
    printf("train_ratio: %f\n", config->train_ratio);
    printf("k: %d\n", config->k);
    printf("nprobe: %d\n", config->nprobe);
    printf("R: %d\n", config->R);
    printf("L: %d\n", config->L);
    printf("B: %d\n", config->B);
    printf("M: %d\n", config->M);
    printf("n: %d\n", config->n);
    printf("dim: %d\n", config->dim);
    printf("max_query: %d\n", config->max_query);
    printf("thread_count: %d\n", config->thread_count);
    printf("recall_k: %d\n", config->recall_k);
    printf("real_k: %d\n", config->real_k);
    printf("last_layer_nprobe: %d\n", config->last_layer_nprobe);
    printf("==========================================\n");
}

Bin* init_bin(const char *filename, int feature_size){
    Bin* bin = (Bin*)malloc(sizeof(Bin));
    bin->size = feature_size;

    //首先读取文件头，获取向量数和向量维数
    FILE *fp = fopen(filename, "rb");
    if (fread(&bin->n, sizeof(int), 1, fp) != 1) {
        perror("Failed to read file");
        exit(EXIT_FAILURE);
    }
    if (fread(&bin->dim, sizeof(int), 1, fp) != 1) {
        perror("Failed to read file");
        exit(EXIT_FAILURE);
    }

    int feature_bytes_count = bin->n * bin->dim *  bin->size;
    //获取文件大小
    fseek(fp, 0, SEEK_END);
    long file_size = ftell(fp);
    fseek(fp, 8, SEEK_SET);
    //检查文件大小是否和feature_bytes_count一致
    if (file_size != feature_bytes_count + 8) {
        fprintf(stderr, "File (%s) size does not match feature size in head(%d)\n",filename, feature_bytes_count+8);
        exit(EXIT_FAILURE);
    }

    //为bin->features分配内存
    bin->features = (void *)malloc(feature_bytes_count);
    //读取文件内容到bin->features中
    if (fread(bin->features, 1, feature_bytes_count, fp) != feature_bytes_count) {
        perror("Failed to read file");
        exit(EXIT_FAILURE);
    }

    return bin;
}

void destroy_bin(Bin* bin){
    free(bin->features);
    free(bin);
}

void* get_bin_line(Bin* bin, int index){
    return (void*)((char*)bin->features + index * bin->size * bin->dim);
}

void print_bin(Bin* bin, int index_start, int index_end, char* format_string){
    if(index_start < 0)
        index_start = 0;
    if(index_end > bin->n)
        index_end = bin->n;
    if(index_start > index_end)
        index_start = index_end;

    printf("Print bin (n=%d, dim=%d, size=%d) from %d to %d: \n",bin->n, bin->dim, bin->size, index_start, index_end);
    for (int i = index_start; i < index_end; i++) {
        printf("[%d]: ",i);
        char* line = (char*)get_bin_line(bin, i);
        for(int j = 0; j < bin->dim; j++){
            char* feature = line + j * bin->size; //指向当前数据的指针

            //根据format_string的类型进行不同的打印方法
            if(strcmp(format_string, "float") == 0){
                float value = *(float*)feature;
                printf("%f ", value);
            }

            if(strcmp(format_string, "uint32") == 0){
                unsigned value = *(unsigned*)feature;
                printf("%u ", value);
            }

            if(strcmp(format_string, "uint8") == 0){
                unsigned char value = *(unsigned char*)feature;
                printf("%u ", value);
            }
        }
        printf("\n");
    }
}

//传入feature_size的原因是需要匹配不同类型的feature
LastLayer_Vectors* create_last_layer_vectors(int count){
    LastLayer_Vectors* vectors = (LastLayer_Vectors*)malloc(sizeof(LastLayer_Vectors));
    vectors->count = count;
    vectors->vectors = (LastLayer_Vector*)malloc(sizeof(LastLayer_Vector) * count);
    return vectors;
}

//一次批量加载多个向量，减少IO次数.
int load_last_layer_vectors(LastLayer_Vectors* vectors, FILE* fp, int index_start, int count){
    ASSERT_NULL_POINTER(vectors, "vectors");
    ASSERT_NULL_POINTER(vectors->vectors, "vectors->vectors");
    ASSERT_NULL_POINTER(fp, "fp");

    if(index_start < 0)
        index_start = 0;
    
    if(count > vectors->count){ //如果count大于最大可存储的向量数量
        printf("Count (%d) is greater than max count (%d), set count to max count\n", count, vectors->count);
        count = vectors->count;
    }

    //计算出读取文件的起始偏移量
    long offset = index_start * sizeof(LastLayer_Vector);
    fseek(fp, offset, SEEK_SET);
    //一次性加载多个向量到缓冲区中
    if(fread(vectors->vectors, sizeof(LastLayer_Vector), count, fp) != count){
        perror("Failed to read file");
        exit(EXIT_FAILURE);
    }
    
    return count;
}

void destroy_last_layer_vectors(LastLayer_Vectors* vectors){
    ASSERT_NULL_POINTER(vectors, "vectors");
    ASSERT_NULL_POINTER(vectors->vectors, "vectors->vectors");

    free(vectors->vectors);
    free(vectors);
}

void print_last_layer_vectors(LastLayer_Vectors* vectors){
    ASSERT_NULL_POINTER(vectors, "vectors");
    ASSERT_NULL_POINTER(vectors->vectors, "vectors->vectors");
    printf("Print last layer vectors (count=%d): \n", vectors->count);
    for(int i = 0; i < vectors->count; i++){
        LastLayer_Vector* vector = &vectors->vectors[i];
        printf("[%d]: ", i);
        printf("id=%d, ", vector->id);
        for(int j = 0; j < 8; j++){
            printf(FEATURE_FORMAT, vector->features[j]);
        }
        printf("\n");
    }
}

//计算两个向量之间的L2距离
float l2_distance_float(float* A, float* B, int dim){
    float result = 0;
	float temp = 0;
	for(int i=0; i<dim; i++){
		temp = A[i]-B[i];
		temp = temp*temp;
		result += temp;
	}
	result = sqrt(result);
	return result;
}

float l2_distance_uint8(unsigned char* A, unsigned char* B, int dim){
    float result = 0;
	unsigned char temp = 0;
	for(int i=0; i<dim; i++){
		temp = A[i]-B[i];
		temp = temp*temp;
		result += temp;
	}
	result = sqrt(result);
	return result;
}

float l2_distance(FEATURE* A, FEATURE* B, int dim){
    float result = 0;
	FEATURE temp = 0;
	for(int i=0; i<dim; i++){
		temp = A[i]-B[i];
		temp = temp*temp;
		result += temp;
	}
	result = sqrt(result);
	return result;
}

int min(int a, int b){
    if(a>b)
        return b;
    return a;
}

float recall_rate(struct minlist* result, Bin* gts, int id, int k){
    //首先确定真实的k值，从k，minlist长度，gts长度中取最小值
    k = min(k, result->count);
    k = min(k, gts->dim);

    float recall = 0.0;
    unsigned int* gt_line = (unsigned int*)get_bin_line(gts, id);

    int same = 0;
    int next_gt = 0;
    for(int i=0; i<k; i++){ //对于向量i的检索结果中的第j个近似向量
        while(next_gt < k){ 
            if(get_minlist_id(result,i) == gt_line[next_gt]){
                same++;
                next_gt++;
                break;
            }
            next_gt++;
        }
        if(next_gt >= k){
            break;
        }
    }
    recall = ((float)same)/k;

    return recall;
}

Request_Results* create_result(){
    Request_Results* ret = (Request_Results*)malloc(sizeof(Request_Results));
    memset(ret, 0, sizeof(Request_Results));
    
    return ret;
}

void print_result(Request_Results* results, Config* config){
    ASSERT_NULL_POINTER(results, "results");
    printf("==================Result==================\n");
    printf("Total %d query\n", results->count);
    printf("Recall rate: %7.4f %%\n", results->recall_rate_sum/results->count);
    printf("Average IO count: %d\n", results->io_count_sum/results->count);
    printf("Average IO length per query: %d, %6.4f %% of total dataset\n", results->io_length_sum/results->count, ((float)(results->io_length_sum/results->count)*100)/config->n);
    printf("Average IO length per IO: %d\n", results->io_length_sum/results->io_count_sum);
    printf("==========================================\n");
}

void insert_result(Request_Results* results, Request_Result* result){
    ASSERT_NULL_POINTER(results, "results");
    ASSERT_NULL_POINTER(result, "result");

    Request_Result* node = (Request_Result*)malloc(sizeof(Request_Result));
    memcpy(node, result, sizeof(Request_Result));
    node->next = NULL;

    //插入节点
    if(results->count == 0){ //如果还没有节点，那么头尾指针都指向新节点
        results->result_head = node;
        results->result_tail = node;
    }
    else{ //如果有节点，那么将新节点插入尾部
        results->result_tail->next = node;
        results->result_tail = node;
    }

    //更新统计值
    results->count++;
    results->io_count_sum += result->io_count;
    results->io_length_sum += result->io_length;
    results->recall_rate_sum += result->recall_rate;
}

void destroy_result(Request_Results* results){
    ASSERT_NULL_POINTER(results, "results");

    while(results->result_head != NULL){
        Request_Result* node = results->result_head;
        results->result_head = results->result_head->next;
        free(node);
    }
    free(results);
}